package com.zx.sa.controller;

import java.util.UUID;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.memory.InMemoryChatMemoryRepository;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Controller;
import org.springframework.ui.Model;
import org.springframework.web.bind.annotation.CrossOrigin;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.ResponseBody;
import reactor.core.publisher.Flux;

@Controller
@CrossOrigin(origins = "*")
@RequestMapping("/advisors2")
public class ChatAdvisorsController2 {

    private final ChatClient chatClient;

    @Autowired
    public ChatAdvisorsController2(ChatModel chatModel) {
        // ChatMemoryRepository 是内存存储会话。
        // 除此之外，还有：JdbcChatMemoryRepository、CassandraChatMemoryRepository、Neo4jChatMemoryRepository
        // 可也可以自定义Redis等工具存储会话
        ChatMemoryRepository repository = new InMemoryChatMemoryRepository();
        // 支持多种内容的存储：
        //   1. MessageChatMemoryAdvisor 存储消息，例如：MessageWindowChatMemory
        //   2. PromptChatMemoryAdvisor  存储提示词
        //   3. VectorStoreChatMemoryAdvisor 存储向量
        ChatMemory chatMemory = MessageWindowChatMemory.builder()
                .chatMemoryRepository(repository)
                .maxMessages(5)
                .build();

        this.chatClient= ChatClient.builder(chatModel)
                // 设置默认的角色
                .defaultSystem("你是张雪的个人生活管家2，回答各种问题！")
                // 在使用 Chat Memory 时，可以实现 Spring AI 存储上下文。
                .defaultAdvisors(MessageChatMemoryAdvisor.builder(chatMemory).build())
                // 实现 Logger 的 Advisor
                .defaultAdvisors(new SimpleLoggerAdvisor())
                .build();
    }

    @RequestMapping("/index")
    public String index(Model model) {
        // 每次打开页面都分配一个uuid作为当前会话的id，每个页面的聊天内容都是独立分开的
        model.addAttribute("aiSessionId", UUID.randomUUID().toString());
        return "advisors2";
    }


    @GetMapping(value = "/chat", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    @ResponseBody
    public Flux<String> chat(String msg, String aiSessionId) {
        // 执行user时，带上aiSessionId，用于找到会话
        return chatClient.prompt()
                .user(msg)
                .advisors(a-> a.param(ChatMemory.CONVERSATION_ID, aiSessionId))
                .stream()
                .content();
    }

} 